import os
import lmdb
import cv2
import numpy as np
from tqdm import tqdm

def imwrite_unicode(filename, img):
    # 支持中文路径保存图片
    ext = os.path.splitext(filename)[1]
    result, buf = cv2.imencode(ext, img)
    if result:
        with open(filename, mode='wb') as f:
            buf.tofile(f)
        return True
    return False

def lmdb_to_images(lmdb_path, output_dir, gt_file):
    if not os.path.isdir(lmdb_path):
        raise FileNotFoundError(f"LMDB directory not found: {lmdb_path}")
    if not (os.path.exists(os.path.join(lmdb_path, "data.mdb")) and os.path.exists(os.path.join(lmdb_path, "lock.mdb"))):
        raise FileNotFoundError(f"LMDB files not found in directory: {lmdb_path}")
    os.makedirs(output_dir, exist_ok=True)
    env = lmdb.open(lmdb_path, readonly=True, lock=False)
    with env.begin() as txn:
        num_samples = int(txn.get('num-samples'.encode()).decode())
        with open(gt_file, 'w', encoding='utf-8') as gt:
            for idx in tqdm(range(1, num_samples + 1)):
                image_key = 'image-%09d'.encode() % idx
                label_key = 'label-%09d'.encode() % idx
                image_bin = txn.get(image_key)
                label = txn.get(label_key)
                if image_bin is None or label is None:
                    continue
                img_buf = np.frombuffer(image_bin, dtype=np.uint8)
                img = cv2.imdecode(img_buf, cv2.IMREAD_UNCHANGED)
                if img is None:
                    print(f"Warning: Failed to decode image {idx}")
                    continue
                img_filename = f'{idx:09d}.jpg'
                img_path = os.path.join(output_dir, img_filename)
                success = imwrite_unicode(img_path, img)
                if not success:
                    print(f"Warning: Failed to save image {img_path}")
                    continue
                gt.write(f'{img_filename}\t{label.decode("utf-8")}\n')

if __name__ == '__main__':
    # 示例用法
    lmdb_path = r'D:\songlin\Python\26_检修后续实验\ABINet\data\evaluation\SVTP'  # 注意：这里直接指向MJ_test目录
    output_dir = r'D:\songlin\Python\26_检修后续实验\ABINet\data\evaluation\SVTP\lmdb_restore'
    gt_file = r'D:\songlin\Python\26_检修后续实验\ABINet\data\evaluation\SVTP\lmdb_restore\gt.txt'
    lmdb_to_images(lmdb_path, output_dir, gt_file)
